import os
import pandas as pd


class NotFoundHdrFlag(Exception):
    def __init__(self, msg):
        self.msg = msg

    def __str__(self):
        return self.msg


def get_excel_files_in_dir(dir_path, ext=('.xls', '.xlsx')):
    '''
    获取指定目录下的所有指定扩展名的文件列表。
    :param dir_path: 要进行遍历的目录。
    :param ext: 指定的扩展名。
    :return: 在指定目录下的所有具有给定扩展名的文件列表。
    '''
    all_path = []
    for root, dirs, files in os.walk(dir_path, topdown=False):
        for file in files:
            full_path = os.path.join(root, file)
            if file[0] == '~':
                continue
            file_ext = os.path.splitext(file)[1]
            if file_ext in ext:
                all_path.append(full_path)
    return all_path


def detect_hdr_callback(hdr_data, uniqe_flag):
    """
    检测Excel头部标志的回调函数，可以用户自己实现。
    :param hdr_data:文件头部的原始数据。
    :param uniqe_flag:用于定位文件头部内容的唯一性标识。
    :return: 标志位置（row,col）,否则为（None，None）
    """
    pos_init = hdr_data[hdr_data == uniqe_flag]
    pos = pos_init.dropna(how='all')
    if pos.empty:
        raise NotFoundHdrFlag("文件头部未发现定位标识：[{}].".format(uniqe_flag))
    else:
        index_str = pos.index.values[0]
        index_i = pos_init.index.tolist().index(index_str)
        columns_notna = pos.dropna(axis=1)
        columns_str = columns_notna.columns.values[0]
        columns_i = pos.columns.tolist().index(columns_str)
        return index_i, columns_i


def get_column_area_callback(hdr_data, hdr_row_i, hdr_col_i):
    """
    根据标识获取有效的数据列名区域。一般需要手动实现。此功能只获取列名数据区域位置，此位置之下为有效数据区域。
    此例中默认实现以给定位置开始，到备注结束的头部区域位置。可用于处理首列以序号或标号等开关，以备注结尾的单行列名区域。
    :param hdr_data:文件头部的原始数据。
    :param hdr_row_i:头部文件标识行索引。
    :param hdr_col_i:头部文件标识列索引。
    :return:列名区域（行起始，行结束，列起始，列结束）
    """
    lst_flag = "备注"
    lst_row_i, lst_col_i = detect_hdr_callback(hdr_data, lst_flag)
    print("获取有效列名区域的标识[{}]的位置：{}".format(lst_flag, (lst_row_i, lst_col_i)))
    return hdr_row_i, lst_row_i, hdr_col_i, lst_col_i


def change_hdr_callback(excel_path, columns, hdr_row_i, lst_row_i, hdr_col_i, lst_col_i, sheet_name=None):
    '''
    用于修改文件的头部内容的回调函数。一般需要根据实际情况进行调整。这里是一个示例。
    :param excel_path: 要修改的文件路径。
    :param columns: 要修改的列名。
    :param sheet_name: 要修改的表名,默认为活动表。
    :param hdr_row_i: 列名区域行起始。
    :param lst_row_i: 列名区域行结束。
    :param hdr_col_i: 列名区域列起始。
    :param lst_col_i: 列名区列行结束。
    :return: 修改后的数据，可以通过to_excel进行输出。
    '''
    # columns_str = ['序号', '单位名称', '身份证号码', '姓名', '银行账号', '应发合计', '应发养老部分待遇金额',
    #                "应发年金部分待遇金额"]
    # data = pd.read_excel(file_path, skiprows=row_i + 2, dtype={2: 'object', 4: 'object'}, header=None,
    #                      usecols=np.arange(0, len(columns_str)))
    # data.columns = columns_str

    # 检查给定的列名是否与列名列起始和结束匹配
    if lst_col_i - hdr_col_i + 1 != len(columns):
        raise ValueError("columns 长度与 hdr_col_i, lst_col_i 不匹配！")
    use_cols = list(range(hdr_col_i, lst_col_i + 1))
    skiprows = lst_row_i + 1
    if sheet_name is None:
        data = pd.read_excel(excel_path, skiprows=skiprows, usecols=use_cols, header=None)
    else:
        data = pd.read_excel(excel_path, sheet_name=sheet_name, skiprows=skiprows, usecols=use_cols, header=None)
    data.columns = columns
    return data


class ChangeExcelHeader():
    """获取并判断Excel头部内容的一致性，修改数据表的头部列名区域为标准格式。

    使用方法：
    1.创建ChangeExcelHeader()对象。
      ceh = ChangeExcelHeader(excel_path)
    2.获取头部定位唯一标识：
      flag_pos = ceh.detectHeaderFlag(uniqe_flag="序号")
    3.获取文件中的头部列名区域:
      col_area_range = ceh.getColumnAreaRange(get_column_area_callback=get_column_area_callback) #此处通过需要传入自定义get_column_area_callback函数用于确认列名区域范围。
    4.获取列名区域内容。用以进行比较头部相似性等内容。实际上常用对其进行封装的equalTest方法进行对象间的比较，但需要先执行此方法。
      注：保留此方法而不并入equalTest方法中自动调用的原因是为了通过此方法检查头部列名区域的获取结果，便于确定问题，后续可视情况合并。
      col_area_data = ceh.getColumnArea()
    5.与其他相似结构文件对象进行相似性检查。主要用于对相似结构的Excel文件进行批量头部列名区域修改前的一致性测试，只有所有待修改的
      文件列名区域都相同时，才可以进行下一步统一修改列名，否则不能进行。
      if ceh.equalTest(other):
        print("same")
      else:
        print("different")
    6.修改文件头部列名区域内部容。若为单个文件，则可直接从第4步跳转至此。若为批量修改，则需要确保5步中的所有文件都满足same条件。
      used_cols = ['序号', "nouse1", '户主姓名', '残疾人姓名', '户籍类别','乡镇', '村', '家庭住址', '残疾证号', '身份证号码', '年龄', '性别', '低保证号', '一账折姓名', '一折统账号', '审批时间', '是否是建档立卡户', '是否低保户', '低保类别', '残疾类别', '残疾等级', "nouse2", '生活（元）', "nouse3", '护理（元）', "nouse4", '叠加两项补贴(元）', '补发', '12月应发放资金', '备注']
      data = ceh.changeHeaderColumns(used_cols,change_hdr_callback=change_hdr_callback)
    """

    def __init__(self, excel_path, sheet_name=None, hdr_row=10):
        """
        创建一个Excel表头部修改对象。可以通过回调函数进行Excel表的头部数据的修改。

        :param excel_path:要处理的Excel文件路径。
        :param sheet_name:要处理的Excel表表名，若为空则表示默认活动的Excel表。
        :param hdr_row:要进行头部检测的行数范围。为了优化处理速度，判定头部时只读取指定行数，判定头部位置在此范围内进行。
        """
        self.excel_path = excel_path
        self.sheet_name = sheet_name
        self.hdr_row = hdr_row
        self.hdr_flag_row_i = None
        self.hdr_flag_col_i = None
        self._column_area_row_start = None
        self._column_area_row_end = None
        self._column_area_col_start = None
        self._column_area_col_end = None
        self.column_area_data = None

        # 为了优化对大数据文件的读写速率，先只读取文件头部数据。
        if sheet_name is None:
            self.hdr_data = pd.read_excel(self.excel_path, header=None, nrows=hdr_row)
        else:
            self.hdr_data = pd.read_excel(self.excel_path, sheet_name=sheet_name, header=None, nrows=hdr_row)

    def detectHeaderFlag(self, uniqe_flag, detect_hdr_callback=detect_hdr_callback):
        """
        检测文件中给定的头部标志的位置。
        :param detect_hdr_callback:用于识别头部标志的算法，可以自己实现。默认实现为查找出头部数据中，给定定位标识的位置。
        :param uniqe_flag: 用于定位文件头部内容的唯一性标识。
        :return:头部唯一标识的位置。

        回调函数的实现参考：
           excelhdrc.detect_hdr_callback
        """
        self.hdr_flag_row_i, self.hdr_flag_col_i = detect_hdr_callback(self.hdr_data, uniqe_flag)
        if self.hdr_flag_row_i is None and self.hdr_flag_col_i is None:
            self.hdr_flag_row_i = 0
            self.hdr_flag_col_i = 0
        return self.hdr_flag_row_i, self.hdr_flag_col_i

    def getColumnAreaRange(self, get_column_area_callback=get_column_area_callback):
        """
        获取文件中的头部列名区域。

        :param get_column_area_callback:根据标识获取列名区域的回调函数。一般需要手动实现。
        :return:文件头部的列名区域：（行起始，行结束，列起始，列结束）

        回调函数的实现参考：
           excelhdrc.get_column_area_callback
        """
        if self.hdr_flag_row_i is None and self.hdr_flag_col_i is None:
            print("请先执行detectHeaderFlag方法获取头部标识的位置！")
            return None, None, None, None
        self._column_area_row_start, self._column_area_row_end, self._column_area_col_start, self._column_area_col_end = get_column_area_callback(
            self.hdr_data, self.hdr_flag_row_i, self.hdr_flag_col_i)
        return self._column_area_row_start, self._column_area_row_end, self._column_area_col_start, self._column_area_col_end

    def getColumnArea(self):
        """
        获取列名区域内容。用以进行比较头部相似性等内容。
        :return:DataFrame:头部列名区域的原始数据。
        """
        if self._column_area_row_start is None:
            print("请先执行getColumnAreaRange方法获取头部列名区域或检查其运行结果！")
            return None
        self.column_area_data = self.hdr_data.iloc[self._column_area_row_start:self._column_area_row_end + 1,
                                self._column_area_col_start:self._column_area_col_end + 1]
        return self.column_area_data

    def equalTest(self, other,map_dict={}):
        """
        判断两个文件的头部数据是否一致。
        :param other:另一个ChangeExcelHeader对象实例。
        :return:True：相等；False：不相等
        """
        if self.column_area_data is None:
            print("请先执行getColumnArea方法获取头部列名区域数据或检查其运行结果！")
            raise ValueError("请先执行getColumnArea方法获取头部列名区域数据或检查其运行结果！")
        if other.column_area_data is None:
            print("作为参数的other对象未执行getColumnArea方法获取头部列名区域！")
            raise ValueError("作为参数的other对象未执行getColumnArea方法获取头部列名区域！！")
        this_column_area_data = self.column_area_data.fillna("noname")
        other_column_area_data = other.column_area_data.fillna("noname")
        if map_dict:
            for key,value in map_dict.items():
                this_column_area_data = this_column_area_data.replace(value,key,regex=True)
                other_column_area_data = other_column_area_data.replace(value,key,regex=True)
        #print(this_column_area_data)
        #print(other_column_area_data)
        temp = this_column_area_data == other_column_area_data
        print(temp)
        if not (all(temp.all(axis="columns").values)):
            return False
        else:
            return True

    def changeHeaderColumns(self, columns, change_hdr_callback=change_hdr_callback):
        """
        修改当前Excel的头部文件的列名部分为指定的列名。
        :param columns: 要新修改的列名。
        :param change_hdr_callback:用于修改列名的回调函数，一般也需要用户自己定义。
        :return: 以新columns标识的DataFrame。
        """
        if self._column_area_row_start is None:
            print("请先执行getColumnAreaRange方法获取头部列名区域或检查其运行结果！")
            return None

        return change_hdr_callback(self.excel_path, columns,
                                   self._column_area_row_start,
                                   self._column_area_row_end,
                                   self._column_area_col_start,
                                   self._column_area_col_end,
                                   sheet_name=self.sheet_name)


def get_column_area_callback_2(hdr_data, hdr_row_i, hdr_col_i):
    """
    根据标识获取有效的数据列名区域。一般需要手动实现。此功能只获取列名数据区域位置，此位置之下为有效数据区域。
    此例中默认实现以给定位置开始，到备注结束的头部区域位置。可用于处理首列以序号或标号等开关，以备注结尾的单行列名区域。
    :param hdr_data:文件头部的原始数据。
    :param hdr_row_i:头部文件标识行索引。
    :param hdr_col_i:头部文件标识列索引。
    :return:列名区域（行起始，行结束，列起始，列结束）
    """
    lst_flag = "年金部分待遇金额"
    lst_row_i, lst_col_i = detect_hdr_callback(hdr_data, lst_flag)
    print("获取有效列名区域的标识[{}]的位置：{}".format(lst_flag, (lst_row_i, lst_col_i)))
    return hdr_row_i, lst_row_i, hdr_col_i, lst_col_i

def test_func3():
    excel_path = "../原始数据/测试数据/机关事业养老保险（张立 17793226778）"
    excel_files = get_excel_files_in_dir(excel_path)
    main_test_file = excel_files[0]
    main_ceh = ChangeExcelHeader(main_test_file)
    main_ceh.detectHeaderFlag(uniqe_flag="序号")
    main_ceh.getColumnAreaRange(get_column_area_callback_2)
    main_column_data = main_ceh.getColumnArea()
    print(main_column_data)
    same_flag = True
    for excel_file in excel_files[1:]:
        file_name = os.path.splitext(os.path.basename(excel_file))[0]
        if int(file_name[:4]) != 2022 or int(file_name[-2:])>11:
            continue
        ceh = ChangeExcelHeader(excel_file)
        ceh.detectHeaderFlag(uniqe_flag="序号")
        ceh.getColumnAreaRange(get_column_area_callback_2)
        column_data = ceh.getColumnArea()
        print(column_data)
        if not ceh.equalTest(main_ceh):
            print("{}文件头部结构不一致！".format(excel_file))
            same_flag = False
            break
    if same_flag:
        for excel_file in excel_files:
            file_name = os.path.splitext(os.path.basename(excel_file))[0]
            if int(file_name[:4]) != 2022 or int(file_name[-2:]) > 11:
                continue
            ceh = ChangeExcelHeader(excel_file)
            ceh.detectHeaderFlag(uniqe_flag="序号")
            ceh.getColumnAreaRange(get_column_area_callback_2)
            used_cols = ["序号","单位名称","身份证号码","姓名","银行账号","合计","养老部分待遇金额","年金部分待遇金额"]
            data = ceh.changeHeaderColumns(used_cols)
            print(data.head())


def test_func2():
    '''
    测试批量修改。
    :return:
    '''
    excel_path = "../原始数据/残疾人两项补贴花名册/2022年残疾人两项补贴花名册/"
    excel_files = get_excel_files_in_dir(excel_path)
    main_test_file = excel_files[0]
    main_ceh = ChangeExcelHeader(main_test_file)
    main_ceh.detectHeaderFlag(uniqe_flag="序号")
    main_ceh.getColumnAreaRange()
    main_column_data = main_ceh.getColumnArea()
    same_flag = True
    map_dict = {
        '月应发放资金':"\d*月应发放资金"
    }
    for excel_file in excel_files[1:]:
        print(excel_file)
        ceh = ChangeExcelHeader(excel_file)
        ceh.detectHeaderFlag(uniqe_flag="序号")
        ceh.getColumnAreaRange()
        column_data = ceh.getColumnArea()
        print(column_data)
        if not ceh.equalTest(main_ceh,map_dict=map_dict):
            print("{}文件头部结构不一致！".format(excel_file))
            same_flag = False
            break
    if same_flag:
        for excel_file in excel_files:
            ceh = ChangeExcelHeader(excel_file)
            ceh.detectHeaderFlag(uniqe_flag="序号")
            ceh.getColumnAreaRange()
            used_cols = ['序号', "nouse1", '户主姓名', '残疾人姓名', '户籍类别', '乡镇', '村', '家庭住址', '残疾证号',
                     '身份证号码', '年龄', '性别', '低保证号', '一账折姓名', '一折统账号', '审批时间',
                     '是否是建档立卡户', '是否低保户', '低保类别', '残疾类别', '残疾等级', "nouse2", '生活（元）',
                     "nouse3", '护理（元）', "nouse4", '叠加两项补贴(元）', '补发', '12月应发放资金', '备注']
            data = ceh.changeHeaderColumns(used_cols)
            print(data.head())

def test_func1():
    excel_file_path = "../原始数据/残疾人两项补贴花名册/2022年残疾人两项补贴花名册/12月残疾人两项补贴花名册.xlsx"

    ceh = ChangeExcelHeader(excel_path=excel_file_path)
    flag_pos = ceh.detectHeaderFlag(uniqe_flag="序号")
    print("文件头部标识位置：{}".format(flag_pos))
    col_area_range = ceh.getColumnAreaRange()
    print(col_area_range)
    col_area_data = ceh.getColumnArea()
    print(col_area_data)
    print(col_area_data.iloc[0,:].tolist())
    used_cols = ['序号', "nouse1", '户主姓名', '残疾人姓名', '户籍类别','乡镇', '村', '家庭住址', '残疾证号', '身份证号码', '年龄', '性别', '低保证号', '一账折姓名', '一折统账号', '审批时间', '是否是建档立卡户', '是否低保户', '低保类别', '残疾类别', '残疾等级', "nouse2", '生活（元）', "nouse3", '护理（元）', "nouse4", '叠加两项补贴(元）', '补发', '12月应发放资金', '备注']
    data = ceh.changeHeaderColumns(used_cols)
    print(data)


if __name__ == "__main__":
    test_func3()
